import glob
import os
import sys

import numexpr
import numpy
import tables

from charmicat import mp
from charmicat import report
from charmicat import utils
from charmicat import cutsets

class FilterMonitor(mp.MPMonitor):
    def __init__(self, nproc, nfiles, cutdir, cut, logdir):
        super(FilterMonitor, self).__init__(nproc, logdir=logdir, interval=2)
        self.cut    = cut
        self.cutdir = cutdir
        self.files_tot = nfiles

        self.log_glob_pattern = '*{0}*.out'.format(cut.name)
        self.done_glob_pattern = '*{0}*.done'.format(cut.name)

    def process_logs(self):
        self.files_processed = 0
        self.events_passed   = 0
        self.events_total    = 0

        for path in self.logfiles:
            fp = open(path)
            try:
                nfile, npass, ntot = fp.readlines()[-1].split()
            except IndexError:
                nfile, npass, ntot = 0, 0, 0
            except ValueError:
                nfile, npass, ntot = 0, 0, 0

            self.files_processed += int(nfile)
            self.events_passed   += int(npass)
            self.events_total    += int(ntot)
            
        sys.stdout.write("  {0}/{1} files ({2} events) processed (polls={3})\r".format(self.files_processed, 
                                                                                       self.files_tot,
                                                                                       self.events_total, 
                                                                                       self.npolls))
        sys.stdout.flush()
        
    def finalize(self):
        print
        if hasattr(self.cut, 'pos_str'):
            pos_str = self.cut.pos_str
        else:
            pos_str = self.cut.name


        report.write(os.path.join(self.cutdir, 'report.txt'),
                     self.events_passed, 
                     self.events_total,
                     self.cut.name,
                     pos_str)        


class FilterWorker(mp.MPWorker):
    def __init__(self, logdir, args):
        super(FilterWorker, self).__init__(logdir, args)
        
        self.npass = 0
        self.ntot  = 0
        self.nfile = 0

        self.logfile_tag = args[0].name
        self.errfile_tag = args[0].name
        self.donefile_tag = args[0].name

    def work(self, cut, input_files, opts):     
        self.logfile_tag = cut.name
        self.npass = 0
        self.ntot  = 0
        self.nfile = 0
        for ip, path in enumerate(input_files):
            self.current_input_file = path
            
            cut.process(path)
            
            self.nfile += 1

            self.ntot += cut.events_processed

            if hasattr(cut, 'events_passed'):
                self.npass += cut.events_passed
            else:
                self.npass += cut.events_processed

            self.log(self.nfile, self.npass, self.ntot)
        
    def completed(self):
        with open(os.path.join(self.logdir, '{0}.{1}.done'.format(os.getpid(), self.donefile_tag)), 'w') as fp:
            fp.write("{0} {1} {2}\n".format(self.nfile, self.npass, self.ntot))
                

    def error(self, *args):
        args = list(args)
        args.append("(occured on file `{0}')".format(self.current_input_file))

        super(FilterWorker, self).error(*args)


class Filter(object):
    def __init__(self, name, desc=None):
        self.name = name

        self.desc = desc

        self.update     = None
        self.base_dir   = None
        self.output_dir = None

        self.events_processed = 0
        
        try:
            self.compress   = tables.Filters(complevel=1, complib='lzo', shuffle=True)
        except:
            self.compress   = tables.Filters(complevel=1, complib='zlib', shuffle=True)

            

    def prepare(self, output_dir, opts):
        self.base_dir   = output_dir
        self.output_dir = os.path.join(self.base_dir, "pass")

        if not os.access(self.output_dir, os.F_OK):
            os.mkdir(self.output_dir)
        
        if hasattr(opts, 'complevel'):
            try:
                self.compress = tables.Filters(complevel=opts.complevel, complib='lzo', shuffle=True)
            except ValueError:
                self.compress = tables.Filters(complevel=opts.complevel, complib='zlib', shuffle=True)


    def process_stub(self, input_file):

        path = os.path.join(self.output_dir, os.path.basename(input_file))
        open(path, 'w').close()
        return path

    def process(self, input_file):
        if input_file[-4:] == 'stub':
            return self.process_stub(input_file)

        self.output_files = []
        retval = self._process(input_file)

        fpin = tables.openFile(input_file)

        self.events_processed = 0
        for path in self.output_files:
            fp = tables.openFile(path, 'r+')
            
            for name in fpin.root.nt.attrs._f_list():
                fp.root.nt.attrs[name] = fpin.root.nt.attrs[name]

            if not hasattr(fp.root.nt.attrs, 'filters_applied'):
                fp.root.nt.attrs.filters_applied = [self.name]                
            else:
                fp.root.nt.attrs.filters_applied += [self.name]

            self.events_processed += len(fp.root.nt)

            fp.close()

        fpin.close()
        
        return retval

    def _process(self, input_file):
        tables.copyFile(input_file, self.output_dir)
        return os.path.join(self.output_dir, input_file)

    def finalize(self):
        return glob.glob(os.path.join(self.output_dir, '*'))


        
        

class CutUndefinedAttributeError(Exception):
    """
    An exception class used to stop processing if a cut is
    somehow left undefined. 

    At the moment these likely signify a bug in the code.
    """

class Cut(Filter):
    """
    This class represents a cut to be applied to a PyTables 
    ``Table`` object.

    A ``Cut`` is a collection of methods and attributes to help 
    easily apply, combine, and draw the results of selections one might
    apply to his or her data. 
    
    Instance variables
    ------------------

    name
        The name of the cut, typically identical to the name of the
        associated column of your custom PyTables ``IsDescription`` 
        object.

    pos_str
        A string defining *in PyTables in-kernel query syntax* a 
        conditional statement which evaluates to true for the rows
        to be kept.
    fcn_str
        A Python-compilable string that calculates the value to be cut on.
 
    test_str
        A string defining a Python-compilable equivalent to ``pos_str``
        
    shape
        for multidimensional cuts, the Numpy shape tuple describing the output of fcn_str

    anticut
        specify the inverse of the cut manually with a PyTables-compilable string

    desc
        a more verbose description of the purpose of the cut

    latex
        a LaTeX-compilable representation of this cut

    Public methods
    --------------

    * anticut()
    * eval(r)
    * test(r)

    """
    def __init__(self, name, pos_str, fcn_str='-1', test_str = 'True', shape=(0,), anticut=None, desc=None,
                 downstream_index=None, latex='', condvars={}):
        """
        Create a ``Cut`` instance. 
        
        Arguments:

        name -- The name of the cut, typically identical to the name of the
            associated column of your custom PyTables ``IsDescription`` 
            object.

        pos_str -- A string defining in PyTables in-kernel query syntax a 
            conditional statement which evaluates to true for the events
            you want to keep. For example, a cut on a column ``mm`` to
            keep only events with mm > 0.78 and mm < 1.04 would be ::
                "(mm > 0.78) & (mm < 1.04)"
            Keep in mind that PyTables in-kernel queries use a syntax similar
            but not identical to traditional Python boolean statements. Most 
            importantly, the usual Python constructs ``and`` and ``or`` do *not*
            work as one might expect. Replace these operators with their bitwise
            cousins ``&`` and ``|``, like the above example.

        Keyword arguments:
        
        anticut
            specify the inverse of the cut manually with a PyTables-compilable string

        desc
            a more verbose description of the purpose of the cut


        """
        super(Cut, self).__init__(name, desc) 

        self.pos_str    = pos_str
        """
        A string defining rows which pass the cut 
        (PyTables syntax).
        """

        self._anticut   = anticut
        self.update     = False
        self.fail_dir   = None

        self.downstream_pos_str = None

        self.condvars = condvars

        self.latex = latex

    def _build_downstream_index(self, nt):
        cols_needed = tables._requiredExprVars(self.downstream_pos_str, None, depth=2).keys()
        for col in cols_needed:
            nt.colinstances[col].createIndex()

    def prepare(self, output_dir, opts):
        super(Cut, self).prepare(output_dir, opts)

        self.fail_dir   = os.path.join(self.base_dir, "fail")

        if not os.access(self.fail_dir, os.F_OK):
            os.mkdir(self.fail_dir)

    def process_stub(self, path):
        return super(Cut, self).process_stub(path)
    
    def _process(self, path_in):
        outfile = os.path.basename(path_in)
        path_pass = os.path.join(self.output_dir, outfile)
        path_fail = os.path.join(self.fail_dir, outfile)

        self.output_files = [path_pass, path_fail]

        fppass  = tables.openFile(path_pass, 'w',
                                  "Events which satisify `{0}'".format(self.pos_str),
                                  filters=self.compress)
        fpfail  = tables.openFile(path_fail, 'w',
                                  "Events which do not satisify `{0}'".format(self.pos_str),
                                  filters=self.compress)

        fpin = tables.openFile(path_in)

        
        ntpass = fppass.createTable(fppass.root, 'nt', fpin.root.nt.coldescrs,
                                    "table with events that satisfy `{0}'".format(self.pos_str),
                                    filters=self.compress
                                    )
        ntfail = fpfail.createTable(fpfail.root, 'nt', fpin.root.nt.coldescrs, 
                                    "table with events that satisfy `{0}'".format(self.pos_str),
                                    filters=self.compress
                                    )                                            

        sample_npass = 0
        if len(fpin.root.nt) > 100000:
            sample_npass = len(fpin.root.nt.getWhereList(self.pos_str, stop=10000, condvars=self.condvars))

        # Querying is faster for more selective criteria; use either the default cut
        # or its inverse depending on which selects less events.
        if sample_npass > 5000:
            primary_cut  = self.anticut()
            primary_nt   = ntfail
            secondary_nt = ntpass
        else:
            primary_cut  = self.pos_str
            primary_nt   = ntpass
            secondary_nt = ntfail
        
        step = 500000
        tmp = range(0, len(fpin.root.nt), step)
        tmp.append(len(fpin.root.nt))
        ranges = [(tmp[i], tmp[i+1]) for i in xrange(len(tmp)-1)]

        for low, high in ranges:
            primary_indices = fpin.root.nt.getWhereList(primary_cut, start=low, stop=high, condvars=self.condvars)
            secondary_indices = numpy.setdiff1d(numpy.arange(low, high), primary_indices, assume_unique=True)

            primary_nt.append(fpin.root.nt.readCoordinates(primary_indices))
            secondary_nt.append(fpin.root.nt.readCoordinates(secondary_indices))
            
        npass = len(ntpass)
        nfail = len(ntfail)

        if self.downstream_pos_str is not None:
            self._build_downstream_index(ntpass)
        
        fppass.close()
        fpfail.close()
        fpin.close()
    
        if npass == 0:
            os.remove(path_pass)
            self.output_files.remove(path_pass)
            path_pass = path_pass[:-3] + '.stub'
            open(path_pass, 'w').close()
        elif nfail == 0:
            os.remove(path_fail)
            self.output_files.remove(path_fail)
            path_fail = path_fail[:-3] + '.stub'
            open(path_fail, 'w').close()

                     
        self.events_passed = npass
        self.events_failed = nfail
    
        return path_pass

    def finalize(self):        
        return glob.glob(os.path.join(self.output_dir, '*'))

                                
    def anticut(self):
        """
        Return the cut which selects the opposite of self.pos_str.
        """
        if not self._anticut:
            return '~(' + self.pos_str + ')'
        return self._anticut


    def test(self, rows):                
        """
        Given a PyTables ``Row`` or equivalent, return True if this row would
        pass ``test_str``, and False otherwise.

        Arguments
        ---------
        `rows` -- a Numpy recarray object or equivalent
        """

        rdict = dict([(col, rows[col]) for col in rows.dtype.names])
        rdict.update(self.condvars)

        return numexpr.evaluate(self.pos_str, local_dict=rdict)

    def __str__(self):
        s = "Cut `{0}': `{1}' (`{2}')".format(self.name, self.pos_str, self.desc)

        return s
